package com.ls.rocketmq.agent;

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.utility.JavaModule;
import org.apache.rocketmq.common.protocol.route.BrokerData;
import org.apache.rocketmq.common.protocol.route.TopicRouteData;

import java.security.ProtectionDomain;
import java.util.Map;

public class MQClientAPIImplTransformer implements AgentBuilder.Transformer {

    @Override
    public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder,
                                            TypeDescription typeDescription,
                                            ClassLoader classLoader,
                                            JavaModule module,
                                            ProtectionDomain protectionDomain) {
        return builder
            .method(ElementMatchers.named("getTopicRouteInfoFromNameServer")
                .and(ElementMatchers.takesArguments(String.class, long.class, boolean.class)))
            .intercept(Advice.to(IpReplaceAdvice.class));
    }

    public static class IpReplaceAdvice {

        @Advice.OnMethodExit
        public static void methodEnd(@Advice.Return(readOnly=false) TopicRouteData routeData) {
            try {
                Map<String, String> ipMapping = RocketMqTopicIpReplacerAgent.IP_MAPPING;
//                System.out.println("进入方法返回修改");
//                System.out.println(routeData);
                for (BrokerData brokerData : routeData.getBrokerDatas()) {
                    for (Map.Entry<Long, String> entry : brokerData.getBrokerAddrs().entrySet()) {
                        String[] parts = entry.getValue().split(":");
                        String newIp = ipMapping.get(parts[0]);
                        if (parts.length == 2 && newIp != null) {
                            String newAddr =  newIp + ":" + parts[1];
                            entry.setValue(newAddr);
//                            System.out.println("[RocketMQ Topic Ip Replacer Agent] Replaced IP: " + parts[0] + " -> " + newIp);
                        }
                    }
                }

            } catch (Exception e) {
                e.printStackTrace();
                // 静默处理，不影响主流程
            }
        }

    }
}
